# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from coin.feed.fastfeed import feed_pb2
from coin.strategy.mm.subscription import SubscriptionRequest
from experimental.prophet.graph import graph as graph_mod, graph_def_pb2
from experimental.prophet.graph.shape import Shape


def _ensure_timestamp_placeholder():
  from experimental.prophet.ops import globals
  globals.global_timestamp()


def _register_product(product, graph=None):
  assert not isinstance(product, str)
  symbol_to_product_dict = graph_mod.get_aux_entry('symbol_to_product', graph=graph)
  symbol = str(product)
  if symbol not in symbol_to_product_dict:
    symbol_to_product_dict[symbol] = product
  else:
    assert symbol_to_product_dict[symbol] == product


def find_product_from_graph(symbol, graph=None):
  graph = graph or graph_mod.get_default_graph()
  symbol_to_product_dict = graph.building_ctx.aux_dict.get('symbol_to_product', {})
  return symbol_to_product_dict[symbol]


class _FastFeedPlaceholder:
  def __init__(self, product, sub_req, machine=None, worker_id=None):
    if isinstance(sub_req, SubscriptionRequest):
      sub_req = sub_req.to_str()

    self._attrs = {
        'dtype': graph_def_pb2.DATA_TYPE_FASTFEED,
        'shape': Shape.Scalar().tuple,
        'sub_req': str(sub_req),
        'symbol': str(product),
        'py_product_cls': '%s.%s' % (product.__class__.__module__, product.__class__.__name__)
    }

    if machine is not None:
      self._attrs['machine'] = machine
    if worker_id is not None:
      self._attrs['worker_id'] = worker_id

    # To fill: feed_type, view_type
    self._book = None
    self._trade = None
    self._instrument = None
    self._product = product

  @property
  def symbol(self):
    return str(self._product)

  @property
  def book(self):
    _ensure_timestamp_placeholder()
    _register_product(self._product)
    if self._book is None:
      node_name = '_fastfeed/book.%s' % self.symbol
      attrs = dict(feed_type=feed_pb2.FEED_TYPE_BOOK, **self._attrs)
      with graph_mod.control_if(None):
        node = graph_mod.create_node('Placeholder', [], name=node_name, attrs=attrs, namespace='')
      self._book = node.outputs[0]
    return self._book

  @property
  def trade(self):
    _ensure_timestamp_placeholder()
    _register_product(self._product)
    if self._trade is None:
      node_name = '_fastfeed/trade.%s' % self.symbol
      attrs = dict(feed_type=feed_pb2.FEED_TYPE_TRADE, **self._attrs)
      with graph_mod.control_if(None):
        node = graph_mod.create_node('Placeholder', [], name=node_name, attrs=attrs, namespace='')
      self._trade = node.outputs[0]
    return self._trade

  @property
  def instrument(self):
    _ensure_timestamp_placeholder()
    _register_product(self._product)
    if self._instrument is None:
      node_name = '_fastfeed/instrument.%s' % self.symbol

      feed_type = {
          'Bitmex': 'bitmex_instrument',
          'Okex': 'okex_instrument',
      }[self._product.exchange]

      attrs = dict(feed_type=feed_pb2.FEED_TYPE_CUSTOM,
                   custom_feed_type=feed_type,
                   view_type='json',
                   **self._attrs)
      with graph_mod.control_if(None):
        node = graph_mod.create_node('Placeholder', [], name=node_name, attrs=attrs, namespace='')
      self._instrument = node.outputs[0]
    return self._instrument


def fastfeed_coin(product, sub_req, machine=None, worker_id=None):
  fastfeed_dict = graph_mod.get_aux_entry('fastfeed_coin')
  key = (product, sub_req, machine, worker_id)
  if key not in fastfeed_dict:
    fastfeed_dict[key] = _FastFeedPlaceholder(product, sub_req, machine, worker_id)
  return fastfeed_dict[key]
